"""
模型管理组件
"""

import tkinter as tk
from tkinter import filedialog, messagebox
import threading
from pathlib import Path
from typing import TYPE_CHECKING, Optional
import numpy as np

try:
    from tkinterdnd2 import DND_FILES
    DRAG_DROP_AVAILABLE = True
except ImportError:
    DRAG_DROP_AVAILABLE = False

if TYPE_CHECKING:
    from ..main_window import MainWindow


class ModelManagerMixin:
    """模型管理 Mixin 类"""
    
    def load_model(self: 'MainWindow'):
        """加载 SAM 模型（手动选择）"""
        models_dir = self.config.get('model.models_dir', 'models')
        
        checkpoint_path = filedialog.askopenfilename(
            title="选择 SAM 模型文件",
            initialdir=models_dir if Path(models_dir).exists() else None,
            filetypes=[
                ("模型文件", "*.pth"),
                ("所有文件", "*.*")
            ]
        )
        
        if not checkpoint_path:
            return
        
        # 在后台线程加载
        self.update_status("正在加载模型...")
        self.progress.pack(side=tk.BOTTOM, fill=tk.X, before=self.status_bar)
        self.progress.start()
        
        def load_thread():
            # 使用自动识别模型类型
            auto_detect = self.config.get('model.auto_detect_type', True)
            model_type = None if auto_detect else self.model_type_var.get()
            
            success = self.sam_model.load_model(checkpoint_path, model_type)
            
            self.root.after(0, lambda: self.on_model_loaded(success))
        
        threading.Thread(target=load_thread, daemon=True).start()
    
    def _auto_load_model(self: 'MainWindow'):
        """自动加载最佳可用模型"""
        # 检查是否已有模型加载
        if self.sam_model.is_loaded():
            return
        
        # 查找可用模型
        models = self.sam_model.find_models_in_directory()
        
        if not models:
            self.update_status("未找到模型文件，请手动加载或下载模型")
            self.logger.info("未找到模型文件，跳过自动加载")
            return
        
        # 显示找到的模型
        model_info = ", ".join([f"{m['name']} ({m['type']})" for m in models])
        self.logger.info(f"找到 {len(models)} 个模型: {model_info}")
        
        # 在后台线程自动加载
        self.update_status("正在自动加载模型...")
        self.progress.pack(side=tk.BOTTOM, fill=tk.X, before=self.status_bar)
        self.progress.start()
        
        def load_thread():
            success = self.sam_model.auto_load_best_model()
            self.root.after(0, lambda: self.on_model_loaded(success, auto_loaded=True))
        
        threading.Thread(target=load_thread, daemon=True).start()
    
    def on_model_loaded(self: 'MainWindow', success: bool, auto_loaded: bool = False):
        """
        模型加载完成回调
        
        Args:
            success: 是否成功
            auto_loaded: 是否自动加载
        """
        self.progress.stop()
        self.progress.pack_forget()
        
        if success:
            model_info = self.sam_model.get_model_info()
            
            # 更新模型类型下拉框
            self.model_type_var.set(model_info['model_type'])
            
            status_msg = f"✓ 模型: {model_info['model_type'].upper()} ({model_info['device'].upper()})"
            self.update_status(status_msg)
            
            if auto_loaded:
                # 自动加载成功，只记录日志
                self.logger.info(f"自动加载模型成功: {model_info['model_type']}")
            else:
                # 手动加载成功，显示提示
                messagebox.showinfo(
                    "成功", 
                    f"SAM 模型加载成功！\n\n"
                    f"模型类型: {model_info['model_type'].upper()}\n"
                    f"设备: {model_info['device'].upper()}",
                    parent=self.root
                )
        else:
            self.update_status("模型加载失败")
            if not auto_loaded:
                messagebox.showerror("错误", "加载模型失败！\n请检查模型文件是否正确。", parent=self.root)
    
    def generate_depth_map(self: 'MainWindow'):
        """从当前图片生成景深图"""
        if self.image is None:
            messagebox.showwarning("警告", "请先加载图片！", parent=self.root)
            return
        
        # 检查是否已加载深度估计模型，或模型类型是否匹配
        current_model = getattr(self.depth_estimator, 'model_type', None)
        needs_reload = not self.depth_estimator.is_loaded() or current_model != self.depth_model_type
        
        if needs_reload:
            model_type = self.depth_model_type
            model_sizes = {
                'DPT_Large': '约1.4GB',
                'DPT_Hybrid': '约700MB',
                'MiDaS_small': '约100MB'
            }
            model_size = model_sizes.get(model_type, '未知')
            
            if current_model and current_model != model_type:
                # 模型已加载但类型不同
                msg = f"需要切换到新的深度估计模型。\n\n当前模型: {current_model}\n新模型: {model_type}\n模型大小: {model_size}\n\n是否继续？"
            else:
                # 首次加载
                msg = f"需要加载MiDaS深度估计模型（首次使用会自动下载）。\n\n当前选择: {model_type}\n模型大小: {model_size}\n\n提示：可在工具栏的\"景深模型\"下拉框中切换模型。\n\n是否继续？"
            
            response = messagebox.askyesno("加载深度估计模型", msg, parent=self.root)
            if not response:
                return
            
            # 在后台加载模型
            self.update_status(f"正在加载MiDaS深度估计模型 ({model_type})...")
            self.progress.pack(side=tk.BOTTOM, fill=tk.X, before=self.status_bar)
            self.progress.start()
            
            def load_model_thread():
                # 使用用户选择的模型
                model_type = self.depth_model_type
                self.logger.info(f"加载景深模型: {model_type}")
                success = self.depth_estimator.load_model(model_type)
                self.root.after(0, lambda: self._on_depth_model_loaded(success))
            
            threading.Thread(target=load_model_thread, daemon=True).start()
            return
        
        # 模型已加载且类型正确，直接生成景深图
        self._generate_depth_map_from_image()
    
    def _on_depth_model_loaded(self: 'MainWindow', success: bool):
        """深度估计模型加载完成"""
        self.progress.stop()
        self.progress.pack_forget()
        
        if success:
            self.update_status("MiDaS模型加载成功")
            messagebox.showinfo("成功", "深度估计模型加载成功！\n\n现在可以生成景深图了。", parent=self.root)
            # 自动生成景深图
            self._generate_depth_map_from_image()
        else:
            self.update_status("MiDaS模型加载失败")
            messagebox.showerror("错误", "加载深度估计模型失败！\n\n请检查网络连接后重试。", parent=self.root)
    
    def _generate_depth_map_from_image(self: 'MainWindow'):
        """从当前图片生成景深图（内部方法）"""
        if self.image is None:
            return
        
        self.update_status("正在生成景深图...")
        self.progress.pack(side=tk.BOTTOM, fill=tk.X, before=self.status_bar)
        self.progress.start()
        
        def generate_thread():
            # 生成景深图
            depth_map = self.depth_estimator.estimate_depth(self.image)
            self.root.after(0, lambda: self._on_depth_generated(depth_map))
        
        threading.Thread(target=generate_thread, daemon=True).start()
    
    def _on_depth_generated(self: 'MainWindow', depth_map: Optional[np.ndarray]):
        """景深图生成完成"""
        self.progress.stop()
        self.progress.pack_forget()
        
        if depth_map is None:
            self.update_status("景深图生成失败")
            messagebox.showerror("错误", "生成景深图失败！", parent=self.root)
            return
        
        self.depth_map = depth_map
        self.depth_processor.depth_map = depth_map
        self.depth_processor.set_original_image(self.image)
        
        self.update_status("景深图生成完成")
        
        # 自动开启景深图模式
        self.use_depth_var.set(True)
        self.use_depth_input = True
        
        # 显示景深图
        self.display_depth_map()
        self.left_frame.config(text="景深图（点击选择 - SAM将基于深度信息分割原图）")
        
        self.logger.info("景深图已生成并显示，景深模式已开启")
        
        messagebox.showinfo(
            "成功",
            "景深图生成完成！\n\n"
            "✅ 左侧已显示彩色景深图\n"
            "✅ \"使用景深图\"开关已自动开启\n"
            "✅ SAM将基于深度信息进行精准分割\n\n"
            "🔄 可随时关闭开关切换回普通模式\n\n"
            "现在可以使用单点/框选模式选择目标。",
            parent=self.root
        )
    
    def _on_depth_mode_toggle(self: 'MainWindow'):
        """景深图模式开关切换"""
        self.use_depth_input = self.use_depth_var.get()
        
        if self.use_depth_input:
            # 开启景深模式
            if self.image is None:
                messagebox.showwarning("警告", "请先加载图片！", parent=self.root)
                self.use_depth_var.set(False)
                self.use_depth_input = False
                return
            
            if self.depth_map is None:
                # 还未生成景深图，自动触发生成
                self.logger.info("检测到开启景深模式但未生成景深图，自动触发生成")
                self.generate_depth_map()
                return
            
            # 已有景深图，直接显示
            self.display_depth_map()
            self.update_status("已开启景深模式：SAM将基于景深图进行分割")
            self.left_frame.config(text="景深图（点击选择 - SAM将基于深度信息分割原图）")
            self.logger.info("景深模式已开启：SAM输入=景深图，输出=原图分割结果")
        else:
            # 关闭景深模式，恢复原图
            if self.image is not None:
                self.display_original_image()
                self.update_status("已关闭景深模式：SAM将基于原图进行分割")
                if DRAG_DROP_AVAILABLE:
                    self.left_frame.config(text="原始图片（点击选择对象 / 拖拽图片到此处）")
                else:
                    self.left_frame.config(text="原始图片（点击选择对象）")
                self.logger.info("景深模式已关闭：SAM输入=原图")
    
    def _on_rembg_toggle(self: 'MainWindow'):
        """rembg开关切换"""
        self.use_rembg = self.use_rembg_var.get()
        
        if self.use_rembg:
            # 检查是否有图片
            if self.image is None:
                messagebox.showwarning("警告", "请先加载图片！", parent=self.root)
                self.use_rembg_var.set(False)
                self.use_rembg = False
                return
            
            # 检查rembg是否可用
            if not self.bg_remover.rembg_available:
                messagebox.showerror(
                    "错误",
                    "rembg库未安装！\n\n"
                    "请先安装：\n"
                    "pip install rembg\n\n"
                    "或在下载模型对话框中查看安装说明。",
                    parent=self.root
                )
                self.use_rembg_var.set(False)
                self.use_rembg = False
                return
            
            # 如果已有去背景图片，直接显示
            if self.image_no_bg is not None:
                self._display_rembg_result()
                return
            
            # 执行背景移除
            self.update_status("正在使用Rembg移除背景...")
            self.progress.pack(side=tk.BOTTOM, fill=tk.X, before=self.status_bar)
            self.progress.start()
            
            def remove_bg_thread():
                result = self.bg_remover.remove_background(self.image)
                self.root.after(0, lambda: self._on_rembg_completed(result))
            
            threading.Thread(target=remove_bg_thread, daemon=True).start()
        else:
            # 关闭Rembg，恢复显示
            self.update_status("已关闭Rembg：恢复原图显示")
            self.logger.info("Rembg模式已关闭")
            
            # 根据当前状态决定显示什么
            if self.use_depth_input and self.depth_map is not None:
                self.display_depth_map()
            else:
                self.display_original_image()
    
    def _on_rembg_completed(self: 'MainWindow', result_rgba: Optional[np.ndarray]):
        """Rembg背景移除完成"""
        self.progress.stop()
        self.progress.pack_forget()
        
        if result_rgba is None:
            self.update_status("Rembg背景移除失败")
            messagebox.showerror("错误", "背景移除失败！", parent=self.root)
            self.use_rembg_var.set(False)
            self.use_rembg = False
            return
        
        # 转换RGBA为RGB（白色背景）
        self.image_no_bg = self.bg_remover.get_rgb_from_rgba(result_rgba)
        
        self.logger.info("Rembg背景移除完成")
        
        # 显示去背景后的图片
        self._display_rembg_result()
        
        messagebox.showinfo(
            "成功",
            "背景移除完成！\n\n"
            "✅ 左侧已显示去背景后的图片\n"
            "✅ SAM将基于去背景图进行分割\n\n"
            "注意：首次使用已自动下载模型",
            parent=self.root
        )
    
    def _display_rembg_result(self: 'MainWindow'):
        """显示去背景后的图片"""
        if self.image_no_bg is None:
            return
        
        # 显示去背景图片
        self.display_original_image_custom(self.image_no_bg)
        self.update_status("已显示去背景图片：SAM将基于此图分割")
        self.left_frame.config(text="去背景图片（Rembg处理后）")

